package main

import "net"
import (
	"bufio"
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"
)

var (
	remoteServer = flag.String("h", "192.168.1.49:4398", "the remote server")
	transferData = flag.String("d", "", "the data for sending")
	workerNum    = flag.Int("g", 10, "the number of the workers")
	demoFile     = flag.String("df", "E:\\deps\\test_demangle.cpp", "the file for reading")
)

func convertIP(addr uint32) string {
	dotNum := make([]uint32, 4)
	dotNum[3] = addr & 0xFF
	addr >>= 8
	dotNum[2] = addr & 0xFF
	addr >>= 8
	dotNum[1] = addr & 0xFF
	addr >>= 8
	dotNum[0] = addr & 0xFF
	return fmt.Sprintf("%d.%d.%d.%d", dotNum[0], dotNum[1], dotNum[2], dotNum[3])
}

func handleConn(input *bufio.Reader, output *bufio.Writer) {
	conn, err := net.Dial("tcp", *remoteServer)
	if err != nil {
		log.Printf("failed to connection server :%s", err)
		os.Exit(-1)
	}

	closeSignal := make(chan bool)
	go func() {
		packetSize := int16(0)
		rconn := bufio.NewReader(conn)
		head := make([]byte, 6)
		for {
			n, err := io.ReadAtLeast(rconn, head, len(head))
			if err != nil {
				log.Printf("failed to read head: %s", err)
				break
			}

			if n != len(head) {
				log.Fatalf("failed to read head len(%d)", n)
				break
			}

			//log.Printf("received data from server,len: %d\n", n)
			//log.Println(head)

			r := bytes.NewReader(head)
			var ip uint32 = 0
			binary.Read(r, binary.BigEndian, &packetSize)
			binary.Read(r, binary.BigEndian, &ip)
			log.Printf("got packet size :%d", packetSize)
			//fmt.Fprintf(output, "[%s]:%d-> ", convertIP(ip), packetSize)
			rn, err := io.CopyN(output, rconn, int64(packetSize-4))

			if err != nil {
				log.Printf("failed to copy data :%d", packetSize)
				break
			}
			if rn != int64(packetSize-4) {
				log.Fatalf("unable to read such data from peer:%d,%d", rn, packetSize)
				break
			}

			output.Flush()
		}

		closeSignal <- true

	}()

	for {
		line, err := input.ReadBytes('\n')
		if err != nil {
			log.Printf("failed to read bytes :%s", err)
			break
		}

		//log.Printf("send data to string :%s\n", string(line))
		packetSize := int16(len(line))

		//buf := &bytes.Buffer{}
		//binary.Write(buf, binary.BigEndian, packetSize)
		//
		//sizeBuf := make([]byte, 4)
		//buf.Read(sizeBuf)
		//
		//for i := range sizeBuf {
		//	conn.Write(sizeBuf[i : i+1])
		//	time.Sleep(2 * time.Second)
		//}
		binary.Write(conn, binary.BigEndian, packetSize)
		conn.Write(line)
	}

	time.Sleep(100 * time.Second)
	conn.Close()
	<-closeSignal
}

func main() {
	flag.Parse()

	wg := sync.WaitGroup{}
	wg.Add(*workerNum)
	for i := 0; i < *workerNum; i++ {
		go func(gid int) {
			file, err := os.Open(*demoFile)
			if err != nil {
				log.Fatal(err)
			}
			output, err := os.OpenFile(fmt.Sprintf("received_%d.log", gid), os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
			if err != nil {
				log.Fatal(err)
			}

			handleConn(bufio.NewReader(file), bufio.NewWriter(output))

			output.Close()
			file.Close()
			wg.Done()
		}(i)
	}

	wg.Wait()
}
